import numpy as np
import os
import datetime as dt

def load_data():
    # 加载训练集和测试集数据
    image_log_trainval = np.load('./deconstruction_data/image_log_trainval.npy')
    pv_log_trainval = np.load('./deconstruction_data/pv_log_trainval.npy')
    times_trainval = np.load('./deconstruction_data/times_trainval.npy', allow_pickle=True)

    image_log_test = np.load('./deconstruction_data/image_log_test.npy')
    pv_log_test = np.load('./deconstruction_data/pv_log_test.npy')
    times_test = np.load('./deconstruction_data/times_test.npy', allow_pickle=True)
    
    return image_log_trainval, pv_log_trainval, times_trainval, image_log_test, pv_log_test, times_test

def reconstruct_all_times(times_trainval, times_test):
    # 合并时间戳并按时间排序
    all_times_reconstructed = np.concatenate([times_trainval, times_test])
    all_times_reconstructed.sort()
    return all_times_reconstructed

def find_idx_with_dates(all_times, test_dates):
    idx = []
    for test_day in test_dates:
        test_day_end = test_day + dt.timedelta(days=1)
        # 找到在 [test_day, test_day_end) 范围内的所有时间戳索引
        mask = (all_times >= test_day) & (all_times < test_day_end)
        idx += np.where(mask)[0].tolist()
    return idx

def validate_reconstruction(all_times_reconstructed, test_dates, times_test):
    # 重新计算 idx_test
    idx_test_reconstructed = find_idx_with_dates(all_times_reconstructed, test_dates)
    
    # 提取重建后的测试时间戳
    reconstructed_times_test = all_times_reconstructed[idx_test_reconstructed]
    
    # 验证与保存的 times_test 是否一致
    if not np.array_equal(reconstructed_times_test, times_test):
        raise ValueError("时间戳重建失败，请检查合并或排序逻辑")
    return idx_test_reconstructed

def reconstruct_original_data(
    all_times_reconstructed,
    image_log_trainval,
    pv_log_trainval,
    image_log_test,
    pv_log_test,
    idx_test_reconstructed
):
    # 初始化全量数组
    total_samples = len(all_times_reconstructed)
    
    image_log_reconstructed = np.zeros(
        (total_samples,) + image_log_trainval.shape[1:],
        dtype=image_log_trainval.dtype
    )
    pv_log_reconstructed = np.zeros(
        (total_samples,) + pv_log_trainval.shape[1:],
        dtype=pv_log_trainval.dtype
    )
    
    # 创建 mask_trainval_reconstructed
    mask_trainval_reconstructed = np.ones(total_samples, dtype=bool)
    mask_trainval_reconstructed[idx_test_reconstructed] = False
    
    # 填充数据
    image_log_reconstructed[mask_trainval_reconstructed] = image_log_trainval
    image_log_reconstructed[idx_test_reconstructed] = image_log_test
    
    pv_log_reconstructed[mask_trainval_reconstructed] = pv_log_trainval
    pv_log_reconstructed[idx_test_reconstructed] = pv_log_test
    
    return image_log_reconstructed, pv_log_reconstructed

def restore_original_variables(test_dates):
    # 加载数据
    (image_log_trainval, pv_log_trainval, times_trainval,
     image_log_test, pv_log_test, times_test) = load_data()
    
    # 重建 all_times
    all_times_reconstructed = reconstruct_all_times(times_trainval, times_test)
    
    # 验证并获取 idx_test
    idx_test_reconstructed = validate_reconstruction(
        all_times_reconstructed,
        test_dates,
        times_test
    )
    
    # 重建 image_log 和 pv_log
    image_log, pv_log = reconstruct_original_data(
        all_times_reconstructed,
        image_log_trainval,
        pv_log_trainval,
        image_log_test,
        pv_log_test,
        idx_test_reconstructed
    )
    
    return all_times_reconstructed, image_log, pv_log

# 定义输入参数
sunny_day = [(2017,9,15),(2017,10,6),(2017,10,22),(2018,2,16),(2018,6,12),(2018,6,23),(2019,1,25),(2019,6,23),(2019,7,14),(2019,10,14)]
cloudy_day = [(2017,6,24),(2017,9,20),(2017,10,11),(2018,1,25),(2018,3,9),(2018,10,4),(2019,5,27),(2019,6,28),(2019,8,10),(2019,10,19)]
sunny_datetime = [dt.datetime(day[0],day[1],day[2]) for day in sunny_day]
cloudy_datetime = [dt.datetime(day[0],day[1],day[2]) for day in cloudy_day]
test_dates = sunny_datetime + cloudy_datetime

# 还原变量
all_times, image_log, pv_log = restore_original_variables(test_dates)

np.save('../data_expanded/all_times_highfreq.npy', all_times)
np.save('../data_expanded/all_images_highfreq.npy', image_log)

# 验证结果
print("all_times.shape:", all_times.shape)
print("image_log.shape:", image_log.shape)
print("pv_log.shape:", pv_log.shape)